Continuous-Time Limit of Stochastic Gradient Descent Revisited

Stochastic Gradient Descent (SGD) is an important algorithm in machine learning. With constant learning rates, it is a stochastic process that reaches a stationary distribution. We revisit an analysis of SGD in terms of stochastic differential equations in the limit of small constant gradient steps. This limit, which we feel is not appreciated in the machine learning community, allows us to approximate SGD in terms of a multivariate Ornstein-Uhlenbeck process, and hence to compute stationary distributions in closed form. This formalism has interesting new implications for machine learning. We consider the case where the objective has the interpretation of a log-posterior. Traditional theory suggests choosing the learning rate such that the stationary distribution approximates a point mass at the optimum, but this can lead to wasted effort and overfitting. When the goal is instead to approximate the posterior as well as possible, we can derive criteria for optimal minibatch sizes, learning rates, and preconditioning matrices.